import base64
from openai import OpenAI
from PIL import Image
from config import USER_CONFIGS  # 导入多用户配置
from fastapi import HTTPException

# 初始化客户端，支持多用户配置
def init_client(user_id):
    # 验证用户是否存在
    if user_id not in USER_CONFIGS:
        raise HTTPException(status_code=400, detail=f"用户{user_id}不存在")
    
    # 获取用户特定配置
    user_config = USER_CONFIGS[user_id]['openai']
    return OpenAI(
        base_url=user_config["base_url"],
        api_key=user_config["api_key"]
    )

# 向模型提问并处理图片的函数，增加user_id参数
def ask_model_with_image(system_prompt, user_question, user_id, image_path='./upload/clip_screenshot.png'):
    try:
        # 使用用户特定配置初始化客户端
        client = init_client(user_id)

        # 动态检测图片格式
        with Image.open(image_path) as img:
            image_format = img.format.lower()
            if image_format == 'jpeg':
                mime_type = 'image/jpeg'
            elif image_format == 'png':
                mime_type = 'image/png'
            elif image_format == 'gif':
                mime_type = 'image/gif'
            else:
                raise ValueError(f'不支持的图片格式: {image_format}')

        # 读取图片并转换为 base64
        with open(image_path, "rb") as image_file:
            encoded_image = base64.b64encode(image_file.read()).decode('utf-8')

        # 构造消息
        messages = [
            {
                "role": "system",
                "content": system_prompt
            },
            {
                "role": "user",
                "content": [
                    {
                        "type": "image_url",
                        "image_url": {
                             "url": f'data:{mime_type};base64,{encoded_image}'
                        },
                    },
                    {
                        "type": "text",
                        "text": user_question
                    }
                ]
            }
        ]

        # 调用模型
        response = client.chat.completions.create(
            model="doubao-1.5-vision-lite-250315",
            messages=messages
        )

        return response.choices[0].message.content
    except Exception as e:
        print(f"An error occurred: {e}")
        return None

if __name__ == "__main__":
    system_prompt = "你是一个部落冲突游戏资源识别助手，你只需要返回数字，用逗号隔开，其他什么也不要返回"
    user_question = "金币，圣水，黑油分别为多少"
    image_path = './upload/clip_screenshot.png'
    # 测试时需要指定用户ID
    result = ask_model_with_image(system_prompt, user_question, user_id="user1", image_path=image_path)
    if result:
        print(result)
        # 转换为列表
        result_list = result.split(',')
        print(result_list)